/*
 * Copyright (c) 2024 endless-sky
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "upgrade_impl.h"
#include "mtd_manager.h"
#include "utils_log.h"
#include "common_event.h"
#include <fstream>
#include <filesystem>
#include <sstream>
#include <algorithm>
#include <sys/reboot.h>
#include <unistd.h>
#include <fcntl.h>
#include <cerrno>
#include <nlohmann/json.hpp>

namespace El {
namespace Upgrade {

// 实现IFirmwareUpgradeService接口的工厂函数
IUpgradeService &IUpgradeService::GetInstance()
{
    return UpgradeImpl::GetInstance();
}

UpgradeImpl &UpgradeImpl::GetInstance()
{
    static UpgradeImpl instance;
    return instance;
}

bool UpgradeImpl::Start()
{
    EL_INFO("Starting upgrade service");

    // 标记当前启动为成功状态
    if (!ubootEnvManager_->MarkBootSuccessful()) {
        EL_WARN("Failed to mark boot as successful during startup");
    }

    // 迁移逻辑移除：全新设备无需兼容历史版本
    return true;
}

UpgradeImpl::UpgradeImpl()
    : isRunning_(false),
      upgradeInProgress_(false),
      upgradePhase_(UpgradePhase::WAITING_HEADER),
      cpioParsingState_(CpioParsingState::WAITING_HEADER),
      cpioStreamPosition_(0),
      currentStatus_("IDLE"),
      streamBufferSize_(GetOptimalBufferSize()),
      totalExpectedSize_(0),
      totalActualSize_(0),
      totalWrittenSize_(0),
      progressSizeAdjusted_(false)
{
    ubootEnvManager_ = std::make_unique<UBootEnvManager>();
    mtdManager_ = std::make_unique<MtdManager>();
    ubiManager_ = std::make_unique<UbiManager>();
    // 默认镜像文件名映射
    imageFilenameToType_["kernel.img"] = "kernel";
    imageFilenameToType_["rootfs.img"] = "rootfs";
    imageFilenameToType_["assets.img"] = "assets";
    streamBuffer_.reserve(streamBufferSize_);

    // 初始化CPIO文件信息
    currentFile_ = {};
}

UpgradeImpl::~UpgradeImpl()
{
    AbortUpgrade();
}

bool UpgradeImpl::InitializeUpgradeEnvironment()
{
    if (isRunning_) {
        EL_WARN("Upgrade module already started");
        return true;
    }

    EL_INFO("Preparing MTD-based streaming upgrade module");

    // 初始化MTD设备库
    if (!mtdManager_->InitializeMtdLib()) {
        lastError_ = "Failed to initialize MTD library";
        EL_ERROR("{}", lastError_);
        return false;
    }

    // 验证U-Boot环境配置
    if (!ubootEnvManager_->ValidateBootConfig()) {
        lastError_ = "U-Boot environment validation failed";
        EL_ERROR("{}", lastError_);
        return false;
    }

    // 打印当前分区信息
    ubootEnvManager_->PrintBootInfo();

    // 确定当前活动分区和目标分区
    if (!DetermineTargetPartitions()) {
        lastError_ = "Failed to determine target partitions";
        EL_ERROR("{}", lastError_);
        return false;
    }

    isRunning_ = true;
    currentStatus_ = "READY";
    EL_INFO("MTD streaming upgrade module prepared successfully");
    return true;
}

void UpgradeImpl::AbortUpgrade()
{
    if (!isRunning_) {
        return;
    }

    EL_INFO("Interrupting MTD streaming upgrade module");

    // 取消正在进行的升级
    if (upgradeInProgress_) {
        upgradeInProgress_ = false;
        currentStatus_ = "INTERRUPTED";
    }

    // 清理MTD资源
    if (mtdManager_) {
        mtdManager_->CleanupMtdResources();
    }

    // 清理流式处理状态
    streamBuffer_.clear();
    partitionWritePositions_.clear();
    partitionLogicalWritePositions_.clear();
    partitionWriteInitialized_.clear();
    partitionAlignmentBuffers_.clear();
    
    // 清理进度相关数据
    expectedFileSizes_.clear();
    actualFileSizes_.clear();
    totalExpectedSize_ = 0;
    totalActualSize_ = 0;
    totalWrittenSize_ = 0;
    progressSizeAdjusted_ = false;
    
    upgradePhase_ = UpgradePhase::WAITING_HEADER;

    // 重置CPIO解析状态
    cpioParsingState_ = CpioParsingState::WAITING_HEADER;
    cpioStreamPosition_ = 0;
    currentFile_ = {};

    isRunning_ = false;
    currentStatus_ = "INTERRUPTED";
    EL_INFO("MTD streaming upgrade module interrupted");
}

void UpgradeImpl::RegisterProgressCallback(UpgradeProgressCallback callback)
{
    std::lock_guard<std::mutex> lock(mutex_);
    progressCallback_ = callback;
    EL_INFO("Progress callback set for streaming upgrade");
}

bool UpgradeImpl::StreamUpgradeData(const uint8_t *data, size_t length)
{
    std::lock_guard<std::mutex> lock(mutex_);

    if (!isRunning_) {
        lastError_ = "Upgrade module not prepared";
        EL_ERROR("{}", lastError_);
        return false;
    }

    if (!data || length == 0) {
        lastError_ = "Invalid data for streaming upgrade";
        EL_ERROR("{}", lastError_);
        return false;
    }

    // 第一次接收数据时自动开始升级
    if (!upgradeInProgress_) {
        EL_INFO("Starting streaming upgrade on first data received");
        upgradeInProgress_ = true;
        currentStatus_ = "UPGRADING";
        
        // 重置进度相关变量
        totalWrittenSize_ = 0;
        progressSizeAdjusted_ = false;
        actualFileSizes_.clear();
        
        UpdateProgress(0, "Starting streaming upgrade");
    }

    // 处理流式数据
    if (!ProcessUpgradePackageData(data, length)) {
        upgradeInProgress_ = false;
        currentStatus_ = "FAILED";
        return false;
    }

    return true;
}

bool UpgradeImpl::CompleteUpgrade()
{
    std::lock_guard<std::mutex> lock(mutex_);

    if (!upgradeInProgress_) {
        lastError_ = "No upgrade in progress";
        EL_ERROR("{}", lastError_);
        return false;
    }

    EL_INFO("Finishing streaming upgrade");

    // Update U-Boot environment variables to switch to the new partition
    BootSlot targetSlot = (targetSlot_ == "a") ? BootSlot::SLOT_A : BootSlot::SLOT_B;
    if (!ubootEnvManager_->SwitchBootSlot(targetSlot)) {
        lastError_ = "Failed to switch boot slot to " + targetSlot_;
        EL_ERROR("{}", lastError_);
        upgradeInProgress_ = false;
        currentStatus_ = "FAILED";
        return false;
    }

    // Cleanup streaming state
    streamBuffer_.clear();
    partitionWritePositions_.clear();
    partitionLogicalWritePositions_.clear();
    partitionWriteInitialized_.clear();
    partitionAlignmentBuffers_.clear();
    
    // 清理进度相关数据
    expectedFileSizes_.clear();
    actualFileSizes_.clear();
    totalExpectedSize_ = 0;
    totalActualSize_ = 0;
    totalWrittenSize_ = 0;
    progressSizeAdjusted_ = false;
    
    upgradePhase_ = UpgradePhase::WAITING_HEADER;

    // 重置CPIO解析状态
    cpioParsingState_ = CpioParsingState::WAITING_HEADER;
    cpioStreamPosition_ = 0;
    currentFile_ = {};

    upgradeInProgress_ = false;
    currentStatus_ = "COMPLETED";
    UpdateProgress(100, "Upgrade completed successfully");

    EL_INFO("Streaming upgrade completed successfully. System will boot from slot {} on next reboot", targetSlot_);

    // 启动延迟重启任务，而不是立即重启
    EL_INFO("Scheduling system reboot in 10 seconds to allow web response...");
    UpdateProgress(100, "Upgrade completed, rebooting in 10 seconds...");
    ScheduleDelayedReboot(10);

    return true;
}

std::string UpgradeImpl::GetBootPartitionInfo()
{
    return ubootEnvManager_->GetBootSlotInfo();
}

bool UpgradeImpl::SwitchToBootPartition(const std::string &slot)
{
    BootSlot targetSlot;
    if (slot == "A" || slot == "a") {
        targetSlot = BootSlot::SLOT_A;
    } else if (slot == "B" || slot == "b") {
        targetSlot = BootSlot::SLOT_B;
    } else {
        lastError_ = "Invalid boot slot: " + slot;
        EL_ERROR("{}", lastError_);
        return false;
    }

    return ubootEnvManager_->ForceBootSlot(targetSlot);
}

std::string UpgradeImpl::GetCurrentUpgradeStatus()
{
    std::lock_guard<std::mutex> lock(mutex_);
    return currentStatus_;
}

std::string UpgradeImpl::GetLastErrorMessage()
{
    std::lock_guard<std::mutex> lock(mutex_);
    return lastError_;
}

// ==================== 镜像写入实现 ====================

bool UpgradeImpl::ProcessUpgradePackageData(const uint8_t *data, size_t length)
{
    // 将数据添加到流缓冲区
    streamBuffer_.insert(streamBuffer_.end(), data, data + length);

    // 根据解析状态处理数据
    bool result = false;
    switch (upgradePhase_) {
        case UpgradePhase::WAITING_HEADER:
        case UpgradePhase::PARSING_DESCRIPTION:
        case UpgradePhase::WRITING_IMAGES:
            result = ExtractFilesFromUpgradePackage(streamBuffer_.data(), streamBuffer_.size());
            break;
        default:
            result = false;
    }

    // 如果处理成功，清理已处理的数据
    if (result && cpioStreamPosition_ > 0) {
        // 移除已处理的数据
        if (cpioStreamPosition_ >= streamBuffer_.size()) {
            streamBuffer_.clear();
            cpioStreamPosition_ = 0;
        } else {
            streamBuffer_.erase(streamBuffer_.begin(), streamBuffer_.begin() + cpioStreamPosition_);
            cpioStreamPosition_ = 0;
        }
    }

    return result;
}

bool UpgradeImpl::ExtractFilesFromUpgradePackage(const uint8_t *data, size_t length)
{
    EL_DEBUG("Extracting files from upgrade package: {} bytes, current state: {}", length, static_cast<int>(cpioParsingState_));

    while (cpioStreamPosition_ < length) {
        switch (cpioParsingState_) {
            case CpioParsingState::WAITING_HEADER: {
                // 需要完整的CPIO头部（110字节）
                if (length - cpioStreamPosition_ < sizeof(CpioNewcHeader)) {
                    EL_DEBUG("Not enough data for CPIO header, need {} bytes, have {} bytes", sizeof(CpioNewcHeader),
                              length - cpioStreamPosition_);
                    return true; // 等待更多数据
                }

                // 解析CPIO头部
                if (!ParseCpioHeader(data + cpioStreamPosition_, currentFile_)) {
                    lastError_ = "Failed to parse CPIO header";
                    EL_ERROR("{}", lastError_);
                    return false;
                }

                cpioStreamPosition_ += sizeof(CpioNewcHeader);
                cpioParsingState_ = CpioParsingState::WAITING_FILENAME;

                EL_DEBUG("Parsed CPIO header: filename size={}, file size={}", currentFile_.namesize,
                          currentFile_.filesize);
                break;
            }

            case CpioParsingState::WAITING_FILENAME: {
                // 需要完整的文件名（namesize字节，包含null终止符）
                if (length - cpioStreamPosition_ < currentFile_.namesize) {
                    EL_DEBUG("Not enough data for filename, need {} bytes, have {} bytes", currentFile_.namesize,
                              length - cpioStreamPosition_);
                    return true; // 等待更多数据
                }

                // 提取文件名
                currentFile_.filename = std::string(reinterpret_cast<const char *>(data + cpioStreamPosition_));
                cpioStreamPosition_ += currentFile_.namesize;

                // CPIO要求4字节对齐
                size_t headerAndNameSize = sizeof(CpioNewcHeader) + currentFile_.namesize;
                size_t padding = (4 - (headerAndNameSize % 4)) % 4;
                if (length - cpioStreamPosition_ < padding) {
                    return true; // 等待对齐填充数据
                }
                cpioStreamPosition_ += padding;

                // 记录数据开始位置
                currentFile_.dataOffset = cpioStreamPosition_;
                currentFile_.bytesRead = 0;
                cpioParsingState_ = CpioParsingState::PROCESSING_FILE_DATA;

                EL_INFO("Found file in CPIO: '{}', size: {} bytes", currentFile_.filename, currentFile_.filesize);
                
                // 记录实际文件大小并调整总大小（根据描述映射或默认映射）
                auto it = imageFilenameToType_.find(currentFile_.filename);
                if (it != imageFilenameToType_.end()) {
                    const std::string &imageType = it->second;
                    actualFileSizes_[imageType] = currentFile_.filesize;
                    EL_INFO("Recorded actual file size for {}: {} bytes", imageType, currentFile_.filesize);
                    if (!progressSizeAdjusted_ && actualFileSizes_.size() == expectedFileSizes_.size()) {
                        AdjustProgressTotalSize();
                    }
                }

                // 检查是否是TRAILER表示结束
                if (currentFile_.filename == "TRAILER!!!") {
                    EL_INFO("Reached CPIO archive end");
                    upgradePhase_ = UpgradePhase::WRITING_IMAGES; // 切换到写入完成状态
                    return true;
                }
                break;
            }

            case CpioParsingState::PROCESSING_FILE_DATA: {
                // 处理当前文件的数据
                size_t availableBytes = length - cpioStreamPosition_;
                size_t bytesProcessed = ProcessCurrentFileContent(data + cpioStreamPosition_, availableBytes);

                if (bytesProcessed == 0) {
                    return true; // 需要更多数据或处理失败
                }

                cpioStreamPosition_ += bytesProcessed;
                currentFile_.bytesRead += bytesProcessed;

                // 检查文件是否处理完成
                if (currentFile_.bytesRead >= currentFile_.filesize) {
                    // 文件处理完成，添加数据对齐填充
                    size_t padding = (4 - (currentFile_.filesize % 4)) % 4;
                    if (length - cpioStreamPosition_ < padding) {
                        return true; // 等待填充数据
                    }
                    cpioStreamPosition_ += padding;

                    EL_DEBUG("Completed processing file: '{}'", currentFile_.filename);

                    // 重置状态，准备处理下一个文件
                    cpioParsingState_ = CpioParsingState::WAITING_HEADER;
                    currentFile_ = {};
                }
                break;
            }

            default:
                lastError_ = "Invalid CPIO parsing state";
                EL_ERROR("{}", lastError_);
                return false;
        }
    }

    return true;
}

bool UpgradeImpl::ParseSwDescription(const std::string &content)
{
    try {
        upgradeConfig_ = nlohmann::json::parse(content);
        EL_INFO("Parsed description file successfully");

        // 如果是自研描述格式，建立文件名映射
        try {
            if (upgradeConfig_.contains("format") && upgradeConfig_["format"].is_string() &&
                upgradeConfig_["format"].get<std::string>() == std::string("endless-upgrade")) {
                if (upgradeConfig_.contains("images") && upgradeConfig_["images"].is_object()) {
                    auto images = upgradeConfig_["images"];
                    if (images.contains("kernel") && images["kernel"].is_string()) {
                        imageFilenameToType_[images["kernel"].get<std::string>()] = "kernel";
                    }
                    if (images.contains("rootfs") && images["rootfs"].is_string()) {
                        imageFilenameToType_[images["rootfs"].get<std::string>()] = "rootfs";
                    }
                    if (images.contains("assets") && images["assets"].is_string()) {
                        imageFilenameToType_[images["assets"].get<std::string>()] = "assets";
                    }
                    EL_INFO("Updated image filename mapping from description file");
                }
            }
        } catch (...) {
            EL_WARN("Failed to update image mapping from description file; using defaults");
        }

        // 解析完成后配置分区信息
        if (!ConfigurePartitionsFromSwDescription()) {
            return false;
        }
        descriptionParsed_ = true;

        return true;
    } catch (const std::exception &e) {
        lastError_ = "Failed to parse description file: " + std::string(e.what());
        EL_ERROR("{}", lastError_);
        return false;
    }
}

bool UpgradeImpl::ConfigurePartitionsFromSwDescription()
{
    EL_INFO("Configuring partitions from description for target slot: {}", targetSlot_);

    std::vector<MtdManager::MtdPartitionMap> partitions;

    try {
        // 优先解析自研描述格式
        if (upgradeConfig_.contains("format") && upgradeConfig_["format"].is_string() &&
            upgradeConfig_["format"].get<std::string>() == std::string("endless-upgrade")) {
            // 使用 mtd_map
            if (!upgradeConfig_.contains("mtd_map")) {
                lastError_ = "Invalid description: missing mtd_map";
                EL_ERROR("{}", lastError_);
                return false;
            }
            auto mtd_map = upgradeConfig_["mtd_map"];
            std::string slotKey = (targetSlot_ == "a") ? "slot_a" : "slot_b";
            if (!mtd_map.contains(slotKey)) {
                lastError_ = "Invalid description: missing mtd_map." + slotKey;
                EL_ERROR("{}", lastError_);
                return false;
            }
            auto slot_map = mtd_map[slotKey];
            std::string kernel_mtd = slot_map.contains("kernel") ? slot_map["kernel"].get<std::string>() : "";
            std::string rootfs_mtd = slot_map.contains("rootfs") ? slot_map["rootfs"].get<std::string>() : "";
            std::string assets_mtd = mtd_map.contains("assets") ? mtd_map["assets"].get<std::string>() : "";

            if (kernel_mtd.empty() || rootfs_mtd.empty() || assets_mtd.empty()) {
                lastError_ = "Invalid description: incomplete mtd_map";
                EL_ERROR("{}", lastError_);
                return false;
            }

            partitions.push_back({"kernel", std::string("/dev/") + kernel_mtd, 0, 0, 0, true});
            partitions.push_back({"rootfs", std::string("/dev/") + rootfs_mtd, 0, 0, 0, true});
            partitions.push_back({"assets", std::string("/dev/") + assets_mtd, 0, 0, 0, true});
        } else {
            // 兼容旧 SWUpdate 格式（若存在）
            if (!upgradeConfig_.contains("software") || !upgradeConfig_["software"].contains("stable")) {
                lastError_ = "Invalid sw-description format: missing software.stable section";
                EL_ERROR("{}", lastError_);
                return false;
            }

            auto stable = upgradeConfig_["software"]["stable"];
            std::string targetSlotKey = (targetSlot_ == "a") ? "slota" : "slotb";
            if (!stable.contains(targetSlotKey)) {
                lastError_ = "Target slot configuration not found: " + targetSlotKey;
                EL_ERROR("{}", lastError_);
                return false;
            }
            auto targetSlotConfig = stable[targetSlotKey];
            if (!targetSlotConfig.contains("images")) {
                lastError_ = "No images configuration found for target slot: " + targetSlotKey;
                EL_ERROR("{}", lastError_);
                return false;
            }
            for (const auto &image : targetSlotConfig["images"]) {
                if (image.contains("filename") && image.contains("device") && image.contains("type")) {
                    std::string filename = image["filename"].get<std::string>();
                    std::string device = image["device"].get<std::string>();
                    std::string type = image["type"].get<std::string>();
                    if (type == "flash") {
                        std::string devicePath = "/dev/" + device;
                        std::string imageType = filename;
                        if (imageType.length() > 4 && imageType.substr(imageType.length() - 4) == ".img") {
                            imageType = imageType.substr(0, imageType.length() - 4);
                        }
                        partitions.push_back({imageType, devicePath, 0, 0, 0, true});
                        EL_INFO("Configured target slot {} image: {} -> {} ({})", targetSlot_, filename, imageType,
                                devicePath);
                    }
                }
            }
        }

        if (partitions.empty()) {
            lastError_ = "No valid flash images found for target slot: " + targetSlot_;
            EL_ERROR("{}", lastError_);
            return false;
        }

        // 配置FlashManager的分区信息，FlashManager会从实际设备获取大小信息
        if (!mtdManager_->ConfigurePartitions(partitions)) {
            lastError_ = "Failed to configure partitions in MtdManager: " + mtdManager_->GetLastError();
            EL_ERROR("{}", lastError_);
            return false;
        }
        
        // 计算初始的预期总大小（基于分区大小）
        for (const auto& partition : partitions) {
            auto* mtdPartition = mtdManager_->GetPartition(partition.name);
            if (mtdPartition) {
                totalExpectedSize_ += mtdPartition->size;
                expectedFileSizes_[partition.name] = mtdPartition->size;
                EL_INFO("Initial expected size for {}: {} bytes", partition.name, mtdPartition->size);
            }
        }

        EL_INFO("Successfully configured {} partitions for target slot {}", partitions.size(), targetSlot_);
        return true;

    } catch (const std::exception &e) {
        lastError_ = "Failed to configure partitions from description: " + std::string(e.what());
        EL_ERROR("{}", lastError_);
        return false;
    }
}

bool UpgradeImpl::WriteImageStream(const std::string &imageType,
                                      const uint8_t *data,
                                      size_t length,
                                      bool isLastChunk)
{
    EL_DEBUG("Writing image stream: type={}, length={}, last={}", imageType, length, isLastChunk);

    // 直接使用镜像类型作为分区名称（已经配置为目标slot对应的分区）
    std::string partitionName = imageType;

    // 写入前：若为 assets（静态UBI卷），需要先卸载/分离，避免设备占用
    if (partitionName == "assets") {
        // 仅在第一次写入前做准备
        if (!partitionWriteInitialized_[partitionName]) {
            auto *partition = mtdManager_->GetPartition(partitionName);
            if (!partition) {
                lastError_ = "Partition not found for pre-write: " + partitionName;
                EL_ERROR("{}", lastError_);
                return false;
            }
            // 解析 mtd 设备编号（/dev/mtdX）
            int mtd_num = -1;
            const std::string &dev = partition->device;
            if (dev.rfind("/dev/mtd", 0) == 0) {
                try {
                    mtd_num = std::stoi(dev.substr(8));
                } catch (...) {
                    mtd_num = -1;
                }
            }
            if (mtd_num < 0) {
                lastError_ = "Invalid MTD device path: " + dev;
                EL_ERROR("{}", lastError_);
                return false;
            }
            EL_INFO("Preparing UBI static volume pre-write for assets on {} (mtd{}).", dev, mtd_num);
            if (!ubiManager_->PrepareStaticVolumePreWrite(mtd_num, "/mnt/assets", "assets")) {
                lastError_ = "Failed to prepare assets static volume for write";
                EL_ERROR("{}", lastError_);
                return false;
            }
        }
    }

    // 使用FlashManager写入镜像数据
    if (!mtdManager_->WriteImageStream(partitionName, data, length, isLastChunk, partitionWritePositions_, partitionLogicalWritePositions_,
                                         partitionWriteInitialized_, partitionAlignmentBuffers_)) {
        lastError_ = mtdManager_->GetLastError();
        return false;
    }

    // 获取分区信息以计算进度
    auto *partition = mtdManager_->GetPartition(partitionName);
    if (partition) {
        // 更新已写入的总大小
        totalWrittenSize_ += length;
        
        int progress;
        std::string progressMsg;
        
        if (progressSizeAdjusted_ && totalActualSize_ > 0) {
            // 使用调整后的实际总大小计算进度
            progress = static_cast<int>((totalWrittenSize_ * 100) / totalActualSize_);
            progressMsg = "Writing upgrade files (adjusted size): " + std::to_string(totalWrittenSize_) + "/" + std::to_string(totalActualSize_) + " bytes";
        } else {
            // 使用单个分区的进度（旧逻辑）
            uint64_t logicalOffset = partitionLogicalWritePositions_[partitionName];
            progress = static_cast<int>((logicalOffset * 100) / partition->size);
            progressMsg = "Writing " + imageType + " image to slot " + targetSlot_;
        }
        
        // 确保进度不超过99%（留给完成阶段）
        progress = std::min(progress, 99);
        UpdateProgress(progress, progressMsg);
    }

    // 写入后：若为 assets 的最后一块数据，重新附加并挂载
    if (partitionName == "assets" && isLastChunk) {
        auto *partitionInfo = mtdManager_->GetPartition(partitionName);
        if (!partitionInfo) {
            lastError_ = "Partition not found for post-write: " + partitionName;
            EL_ERROR("{}", lastError_);
            return false;
        }
        int mtd_num = -1;
        const std::string &dev = partitionInfo->device;
        if (dev.rfind("/dev/mtd", 0) == 0) {
            try {
                mtd_num = std::stoi(dev.substr(8));
            } catch (...) {
                mtd_num = -1;
            }
        }
        if (mtd_num < 0) {
            lastError_ = "Invalid MTD device path: " + dev;
            EL_ERROR("{}", lastError_);
            return false;
        }
        EL_INFO("Restoring UBI static volume after assets write on {} (mtd{}).", dev, mtd_num);
        if (!ubiManager_->RestoreStaticVolumePostWrite(mtd_num, "/mnt/assets", "assets")) {
            lastError_ = "Failed to restore/mount assets static volume after write";
            EL_ERROR("{}", lastError_);
            return false;
        }
    }

    return true;
}

bool UpgradeImpl::DetermineTargetPartitions()
{
    // Get current active partition
    BootSlot currentSlot = ubootEnvManager_->GetCurrentBootSlot();
    if (currentSlot == BootSlot::UNKNOWN) {
        lastError_ = "Failed to determine current boot slot";
        EL_ERROR("{}", lastError_);
        return false;
    }

    currentSlot_ = ubootEnvManager_->BootSlotToString(currentSlot);

    // Determine target partition (switch to the other partition)
    targetSlot_ = (currentSlot_ == "a") ? "b" : "a";

    EL_INFO("Current slot: {}, Target slot: {}", currentSlot_, targetSlot_);
    return true;
}

void UpgradeImpl::UpdateProgress(int percentage, const std::string &status)
{
    if (progressCallback_) {
        progressCallback_(percentage, status);
    }

    // 仅首次打印，以及每提升10%时打印一次，减少日志量
    int bucket = percentage / 10; // 0..10
    bool shouldLog = false;
    if (!firstProgressLogged_) {
        shouldLog = true;                // 首次打印
        firstProgressLogged_ = true;
        lastLoggedProgressBucket_ = bucket;
    } else if (bucket > lastLoggedProgressBucket_) {
        shouldLog = true;                // 每跨越一个10%档位打印一次
        lastLoggedProgressBucket_ = bucket;
    }

    if (shouldLog) {
        EL_INFO("Upgrade progress: {}% - {}", percentage, status);
    }
}

void UpgradeImpl::AdjustProgressTotalSize()
{
    totalActualSize_ = 0;
    EL_INFO("Adjusting progress total size based on actual file sizes:");
    
    for (const auto& pair : actualFileSizes_) {
        const std::string& imageType = pair.first;
        uint64_t actualSize = pair.second;
        uint64_t expectedSize = expectedFileSizes_[imageType];
        
        totalActualSize_ += actualSize;
        EL_INFO("  {}: expected {} bytes, actual {} bytes", imageType, expectedSize, actualSize);
    }
    
    EL_INFO("Total size adjusted: expected {} bytes -> actual {} bytes", 
             totalExpectedSize_, totalActualSize_);
    
    progressSizeAdjusted_ = true;
}

size_t UpgradeImpl::GetOptimalBufferSize() const
{
    // 获取可用内存信息
    std::ifstream meminfo("/proc/meminfo");
    std::string line;
    size_t available_memory = 0;

    while (std::getline(meminfo, line)) {
        if (line.find("MemAvailable:") == 0) {
            std::istringstream iss(line);
            std::string key, unit;
            iss >> key >> available_memory >> unit;
            available_memory *= 1024; // 转换为字节
            break;
        }
    }

    // 如果无法获取可用内存，使用默认值
    if (available_memory == 0) {
        EL_WARN("Unable to determine available memory, using default buffer size");
        return DEFAULT_STREAM_BUFFER_SIZE;
    }

    EL_INFO("Available memory: {} MB", available_memory / (1024 * 1024));

    // 根据可用内存动态调整缓冲区大小
    size_t buffer_size;
    if (available_memory >= 8 * 1024 * 1024) {        // >= 8MB
        buffer_size = MAX_STREAM_BUFFER_SIZE;         // 4MB buffer
    } else if (available_memory >= 4 * 1024 * 1024) { // >= 4MB
        buffer_size = DEFAULT_STREAM_BUFFER_SIZE;     // 2MB buffer
    } else {
        buffer_size = MIN_STREAM_BUFFER_SIZE; // 1MB buffer (minimum)
    }

    // 确保缓冲区大小不超过可用内存的40%
    size_t max_allowed = available_memory * 40 / 100;
    buffer_size = std::min(buffer_size, max_allowed);

    // 确保在合理范围内
    buffer_size = std::max(buffer_size, MIN_STREAM_BUFFER_SIZE);
    buffer_size = std::min(buffer_size, MAX_STREAM_BUFFER_SIZE);

    EL_INFO("Optimal stream buffer size: {} MB", buffer_size / (1024 * 1024));
    return buffer_size;
}

bool UpgradeImpl::ParseCpioHeader(const uint8_t *header, CpioFileInfo &fileInfo)
{
    const CpioNewcHeader *cpioHeader = reinterpret_cast<const CpioNewcHeader *>(header);

    // 验证魔法数字 (支持newc和crc格式)
    if (strncmp(cpioHeader->magic, "070701", 6) != 0 && strncmp(cpioHeader->magic, "070702", 6) != 0) {
        EL_ERROR("Invalid CPIO magic number: {:.6s}", cpioHeader->magic);
        return false;
    }

    // 解析文件大小和文件名长度
    fileInfo.filesize = ParseCpioHexField(cpioHeader->filesize, 8);
    fileInfo.namesize = ParseCpioHexField(cpioHeader->namesize, 8);

    // 验证合理的大小限制
    if (fileInfo.namesize > 1024) { // 文件名不应超过1KB
        EL_ERROR("Invalid filename size: {}", fileInfo.namesize);
        return false;
    }

    if (fileInfo.filesize > 100 * 1024 * 1024) { // 单文件不应超过100MB
        EL_ERROR("File too large: {} bytes", fileInfo.filesize);
        return false;
    }

    EL_DEBUG("CPIO header parsed: filesize={}, namesize={}", fileInfo.filesize, fileInfo.namesize);
    return true;
}

uint32_t UpgradeImpl::ParseCpioHexField(const char *hexStr, size_t length)
{
    uint32_t result = 0;
    for (size_t i = 0; i < length; ++i) {
        char c = hexStr[i];
        if (c >= '0' && c <= '9') {
            result = result * 16 + (c - '0');
        } else if (c >= 'a' && c <= 'f') {
            result = result * 16 + (c - 'a' + 10);
        } else if (c >= 'A' && c <= 'F') {
            result = result * 16 + (c - 'A' + 10);
        } else {
            // 忽略非十六进制字符（如null填充）
            break;
        }
    }
    return result;
}

size_t UpgradeImpl::ProcessCurrentFileContent(const uint8_t *data, size_t availableBytes)
{
    if (availableBytes == 0) {
        return 0;
    }

    // 计算需要处理的字节数
    size_t bytesToProcess =
        std::min(availableBytes, static_cast<size_t>(currentFile_.filesize - currentFile_.bytesRead));

    if (bytesToProcess == 0) {
        return 0;
    }

    // 根据文件类型进行处理
    if (currentFile_.filename == "description.json") {
        // 处理描述文件（支持SW兼容或自研格式）
        if (currentFile_.bytesRead == 0) {
            swDescription_.clear();
            swDescription_.reserve(currentFile_.filesize);
        }

        swDescription_.append(reinterpret_cast<const char *>(data), bytesToProcess);

        if (currentFile_.bytesRead + bytesToProcess >= currentFile_.filesize) {
            EL_INFO("description file received: {} bytes from {}", swDescription_.size(), currentFile_.filename);
            if (!ParseSwDescription(swDescription_)) {
                return 0; // 解析失败
            }
            upgradePhase_ = UpgradePhase::PARSING_DESCRIPTION;
        }

    } else if (imageFilenameToType_.find(currentFile_.filename) != imageFilenameToType_.end()) {
        // 处理镜像文件，直接流式写入MTD分区
        std::string imageType = imageFilenameToType_[currentFile_.filename];

        bool isLastChunk = (currentFile_.bytesRead + bytesToProcess >= currentFile_.filesize);
        if (!WriteImageStream(imageType, data, bytesToProcess, isLastChunk)) {
            EL_ERROR("Failed to write image stream for {}", currentFile_.filename);
            return 0;
        }

        upgradePhase_ = UpgradePhase::WRITING_IMAGES;

    } else {
        // 跳过未知文件
        EL_WARN("Skipping unknown file: {}", currentFile_.filename);
    }

    return bytesToProcess;
}

void UpgradeImpl::ScheduleDelayedReboot(int delaySeconds)
{
    // 创建后台线程执行延迟重启
    std::thread rebootThread(&UpgradeImpl::DelayedRebootTask, this, delaySeconds);
    rebootThread.detach(); // 分离线程，让其在后台运行
}

void UpgradeImpl::DelayedRebootTask(int delaySeconds)
{
    EL_INFO("Delayed reboot task started, will reboot in {} seconds", delaySeconds);
    
    // 等待指定的延迟时间
    sleep(delaySeconds);
    
    // 执行系统重启
    EL_INFO("Executing delayed system reboot...");
    if (reboot(RB_AUTOBOOT) < 0) {
        EL_ERROR("Failed to reboot system: {}", strerror(errno));
    }
}

// 默认分区映射函数已移除，统一以描述文件为准。

} // namespace Upgrade
} // namespace El
